import torch
import torch.nn as nn
from perceptron_model import Perceptron
from net_model import Net

# 检查是否有可用的CUDA设备,如果有则使用GPU,否则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

X = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0]).to(device) # 创建输入数据的张量,并将其移动到指定设备上
y = torch.tensor([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0]).to(device) # 创建输入数据的张量,并将其移动到指定设备上

# 单层感知机类模型训练
model1 = Perceptron().to(device) # 创建感知机模型的实例,并将其移动到指定设备上
loss_fnc1 = nn.MSELoss(reduction='mean')
optimizer1 = torch.optim.SGD(model1.parameters(), lr=0.01)

for epoch in range(500):
    y_pred1 = model1(X.unsqueeze(1))
    loss1 = loss_fnc1(y_pred1, y.unsqueeze(1))
    optimizer1.zero_grad()
    loss1.backward()
    optimizer1.step()
    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/500], Loss: {loss1.item():.4f}')

# 保存单层感知机类模型参数
perceptron_path = 'perceptron_model.pth'
torch.save(model1.state_dict(), perceptron_path)
print(f"单层感知机类模型参数已保存到: {perceptron_path}")

# 多层感知机类模型训练
model2 = Net().to(device) # 创建感知机模型的实例,并将其移动到指定设备上
loss_fnc2 = nn.MSELoss(reduction='mean')
optimizer2 = torch.optim.SGD(model2.parameters(), lr=0.01)

for epoch in range(500):
    y_pred2 = model2(X.unsqueeze(1))
    loss2 = loss_fnc2(y_pred2, y.unsqueeze(1))
    optimizer2.zero_grad()
    loss2.backward()
    optimizer2.step()
    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/500], Loss: {loss2.item():.4f}')

# 保存单层感知机类模型参数
net_path = 'net_model.pth'
torch.save(model2.state_dict(), net_path)
print(f"多层感知机类模型参数已保存到: {net_path}")